import cv2 as cv
import numpy as np
from myFFT import myfft_function
import time
from grayToRgb import ima_to_rgb
from videos_to_img import split_video

fft_images = r"E:\studylife\detectflaws\code\python\moNiImg\result3"  # 结果图像的保存路径
videos_path = r"E:\studylife\detectflaws\code\python\moNiImg\videos\new.mp4"  # 原视频存放地址
save_path = r"E:\studylife\detectflaws\code\python\moNiImg\split3"  # 原视频裁切出来的图片的存放地址

# 初始化要计算的图片区域
columns_high = 640
columns_low = 0
rows_high = 480
rows_low = 0
# 下采样次数 Number of downsampling times   1为不做下采样
nodt = 4
# 所计算的视频的起始帧数
map_quantity_begin = 0
# 所计算的视频的结束帧数
map_quantity_end = 300
# 所要选取的频率
Hz = 1
# 是否需要视频转图片
videosToImages = 1

time_start = time.time()  # 开始计时

# 下采样后输出的图片尺寸
row = int(rows_high / nodt)
columns = int(columns_high / nodt)
# 创建输出的幅值列表
list_amp_gray = []
# 创建输出的相位列表
list_pha_gray = []

# 该函数用于读取视频并把视频每帧保存为一张图片
if videosToImages:
    split_video(videos_path, save_path, 1)

# 创建列表存储对应帧数视频图像
lists_gray = []
array_gray = []
for k in range(map_quantity_begin, map_quantity_end):  # 读取map_quantity帧图片计算fft
    img = cv.imread(f'{save_path}/{k}.jpg', flags=0)  # flags=0读取为灰度图像
    # 为列表赋值
    lists_gray.append(img)
array_gray = np.array(lists_gray)  # 把lists_gray转换成np.array类型方便后续裁切

for i in range(rows_low, rows_high, nodt):  # nodt为步长即下采样
    for j in range(columns_low, columns_high, nodt):
        # fft运算得到指定频率幅度值和相位值
        amp_gray, pha_gray = myfft_function(array_gray[:, i, j], map_quantity_begin, map_quantity_end, Hz)
        # print(f'{i}行{j}列, gray:', amp_gray, pha_gray)
        # list_amp_gray.append(amp_gray)
        list_pha_gray.append(pha_gray)

# # 分别处理幅值图和相位图
# # 幅值图
# # print('幅值的列表', list_amp_gray)
# # 将1行的列表重构为图像行列
#
# list_amp_gray_re = np.reshape(list_amp_gray, (row, columns))
# # print('按行列排序后的幅值蓝', list_amp_gray_re)
# # 将列表格式的数据转换为np.array格式便于后续直接输出图像
# array_amp_gray_re = np.array(list_amp_gray_re, dtype=np.uint8)
# array_amp_gray_re = cv.resize(array_amp_gray_re, (640, 480))
# # 保存图片
# cv.imwrite(
#     f'{fft_images}/nodt{nodt},Hz({Hz})amp_gray,r({rows_low},{rows_high})c({columns_low}, {columns_high})m({map_quantity_begin},{map_quantity_end}).jpg',
#     array_amp_gray_re)
# # 展示图像
# cv.imshow('amp_gray', array_amp_gray_re)


# 相位图
# print('相位的列表', list_pha_gray)
# 将1行的列表重构为图像行列
list_pha_gray_re = np.reshape(list_pha_gray, (row, columns))
# print('按行列排序后的相位蓝', list_pha_gray_re)
# 将列表格式的数据转换为np.array格式便于后续直接输出图像
array_pha_gray_re = np.array(list_pha_gray_re, dtype=np.uint8)
array_pha_gray_re = cv.resize(array_pha_gray_re, (640, 480))
# 保存图片
cv.imwrite(
    f'{fft_images}/nodt{nodt},Hz({Hz})pha_gray,r({rows_low},{rows_high})c({columns_low}, {columns_high})m({map_quantity_begin},{map_quantity_end}).jpg',
    array_pha_gray_re)
# 展示图像
cv.imshow('pha_gray', array_pha_gray_re)

# 展示并保存直方图均衡化增强后的相位图
# # 自适应直方图均衡化
# clahe = cv.createCLAHE(clipLimit=300, tileGridSize=(2, 2))  # clipLimit：这是对比度限制的阈值
# image_result = clahe.apply(array_pha_gray_re)  # tileGridSize：将输入图像划分为M × N块，然后对每个局部块应用直方图均衡化

# image_result = cv.equalizeHist(array_pha_gray_re)  # 普通直方图均衡化
# # 保存并展示增强后的图片
# cv.imwrite(
#     f'{fft_images}/plus,nodt{nodt},Hz({Hz})pha_gray,r({rows_low},{rows_high})c({columns_low}, {columns_high})m({map_quantity_begin},{map_quantity_end}).jpg',
#     image_result)
# cv.imshow('Equalized Image', image_result)

# # 使用rgb图显示结果  此方法为使用灰度值彩色对照表实现
# RGBlist=np.loadtxt('ppbyIron.txt')  # 载入色彩对照表
# RGBlist=np.uint8(RGBlist)
# out_img = ima_to_rgb(image_result, RGBlist)  # 灰度转rgb
# # 保存图片并展示图片
# cv.imwrite(
#     f'{fft_images}/rgb,nodt{nodt},Hz({Hz})pha_gray,r({rows_low},{rows_high})c({columns_low}, {columns_high})m({map_quantity_begin},{map_quantity_end}).jpg',
#     out_img)
# cv.imshow('Equalized Image', out_img)

# # 使用rgb图显示结果  此方法为使用opencv中的cv.applyColorMap()函数实现
# # 伪彩色图参数有：cv.COLORMAP_JET  cv.COLORMAP_AUTUMN  cv.COLORMAP_BONE   WINTER   RAINBOW   OCEAN
# rgb = cv.COLORMAP_RAINBOW  # 彩色变换类型
# im_color = cv.applyColorMap(array_pha_gray_re, rgb)
# cv.imwrite(
#     f'{fft_images}/{rgb},nodt{nodt},Hz({Hz})pha_gray,r({rows_low},{rows_high})c({columns_low}, {columns_high})m({map_quantity_begin},{map_quantity_end}).jpg',
#     im_color)
# cv.imshow('RGB Image', im_color)


time_end = time.time()  # 结束计时
time_c = time_end - time_start  # 运行所花时间
print('time cost', time_c, 's')

cv.waitKey(0)
cv.destroyAllWindows()
